from service.das.stock_day_das import StockDateDayDas
from service.das.stock_kline_das import KLineData


class stockAnalyseService:
    def calculate_and_store_kline(self, kline_type: str):
        """
        计算并存储周K、月K、年K数据
        :param kline_type: 'W'、'M' 或 'Y'
        """
        # 从 DAS 获取所有日K数据
        day_data = self.das.get_all_day_data()
        if not day_data:
            print("No day K-line data available.")
            return

        # 转换为 DataFrame
        data = pd.DataFrame([day.__dict__ for day in day_data])

        # 转换日期格式
        data["trade_date"] = pd.to_datetime(data["trade_date"])

        # 根据类型分组
        if kline_type == "W":
            data["group"] = (
                data["trade_date"].dt.to_period("W").apply(lambda r: r.start_time)
            )
        elif kline_type == "M":
            data["group"] = (
                data["trade_date"].dt.to_period("M").apply(lambda r: r.start_time)
            )
        elif kline_type == "Y":
            data["group"] = (
                data["trade_date"].dt.to_period("Y").apply(lambda r: r.start_time)
            )
        else:
            raise ValueError("Invalid kline_type. Must be 'W', 'M', or 'Y'.")

        # 聚合计算
        grouped = data.groupby(["stock_code", "group"])
        aggregated = grouped.agg(
            open=("open", "first"),
            high=("high", "max"),
            low=("low", "min"),
            close=("close", "last"),
            vol=("vol", "sum"),
            amount=("amount", "sum"),
        ).reset_index()
